from torch.utils.data import Dataset
from PIL import Image
from torchvision import transforms
import os
import one_hot

class _dataset(Dataset):
    def __init__(self, root_dir):
        super().__init__()
        self.transforms=transforms.Compose(
            [
                transforms.ToTensor(),
                transforms.Resize((60, 160)),
                transforms.Grayscale()
            ]
        )
        self.image_path = [os.path.join(root_dir, image_name) for image_name in os.listdir(root_dir)]
        # print(self.image_path)

    def __len__(self):
        return self.image_path.__len__()

    def __getitem__(self, index):
        image_path = self.image_path[index]
        image = self.transforms(Image.open(image_path))
        label = image_path.split("/")[-1]
        label = label.split("_")[0]
        label_tensor = one_hot.text2Vec(label)
        label_tensor = label_tensor.view(1, -1)[0] # 拉平
        return image, label_tensor

if __name__ == '__main__':
    train_data = _dataset("./datasets/train/")
    image, label = train_data[0]
    print(label.shape)
    